import re
import argparse
from airsim_wrapper import *
import os
import json
import openai
parser = argparse.ArgumentParser()
parser.add_argument("--prompt", type=str, default="prompts/airsim_basic.txt")
parser.add_argument("--uav_num", type=int, default=3)
parser.add_argument('--local_llm', action="store_true", default=False)
parser.add_argument("--sysprompt", type=str, default="system_prompts/airsim_basic.txt")
args = parser.parse_args()

print(f"Initializing AirSim...")
if args.uav_num == 1:
    client = AirSimWrapper1Uav()
else:
    client = AirSimWrapper3Uav()


class CloudLlm:
    def __init__(self):
        with open("config.json", "r") as f:
            config = json.load(f)

        print("Initializing ChatGPT...")
        openai.api_key = config["OPENAI_API_KEY"]
        # openai.api_base = config["OPENAI_API_BASE"]

        with open(args.sysprompt, "r") as f:
            sysprompt = f.read()

        self.chat_history = [
            {
                "role": "system",
                "content": sysprompt
            },
            {
                "role": "user",
                "content": "move 10 units up"
            },
            {
                "role": "assistant",
                "content": """```python
client.fly_to([client.get_drone_position()[0], client.get_drone_position()[1], client.get_drone_position()[2]+10])
```

This code uses the `moveToPositionAsync()` function to move the drone to a new position that is 10 units up from the current position. It does this by getting the current position of the drone using `get_drone_position()` and then creating a new list with the same X and Y coordinates, but with the Z coordinate increased by 10. The drone will then fly to this new position using `fly_to()`.
                """
            }
        ]


    def ask(self, question):
        self.chat_history.append(
            {
                "role": "user",
                "content": question,
            }
        )
        completion = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=self.chat_history,
            temperature=0
        )
        self.chat_history.append(
            {
                "role": "assistant",
                "content": completion.choices[0].message.content,
            }
        )
        return self.chat_history[-1]["content"]

from swift.llm import get_model_tokenizer, get_template, inference, ModelType, get_default_template_type
class LocalLlm:
    def __init__(self):

        # 新加、、、、、、、、、、、、、、、、、、、、、、、、、、、、、、、
        self.SystemPromots = '''你是一名assistant，帮助我使用无人机的AirSim模拟器。
        当我要求你做某事时，你应该给我使用AirSim完成任务所需的Python代码，然后解释该代码的作用。
        您只能使用我为您定义的函数。
        你不能使用任何其他你认为可能存在的假设函数。
        您可以使用math和numpy等库中的简单Python函数。

        以下我定义的函数，是一些可以用来指挥无人机的功能：
        client.takeoff() - 起飞无人机。
        client.land() - 降落无人机。
        client.get_drone_position() - 将无人机的当前位置返回为与X、Y、Z坐标相对应的3个浮点数。
        client.fly_to([X, Y, Z]) - 将无人机飞行到指定的位置，该位置由对应于X、Y、Z坐标的三个参数组成。
        client.fly_to([0, 0, 0]) - 将无人机回到原点位置
        client.fly_path(points) - 使无人机沿着点数组指定的路径飞行。每个点也是一个由3个浮点数组成的数组，这些浮点数对应于X、Y、Z坐标。
        client.set_yaw(yaw) - 将无人机的偏航设置为以度为单位的指定值。
        client.get_yaw() - 以度为单位返回无人机的当前偏航。
        client.get_position(object_name): 以一个字符串作为输入，指示感兴趣对象的名称，并返回一个由3个浮点数组成的数组，指示其X、Y、Z坐标。
        client.arrage_in_row() - set  the drone in mode 1.
        client.arrage_in_column() - set  the drone in mode 2.
        
        根据轴惯例，正向表示正X轴。右侧表示正Y轴。向下表示正Z轴:
        原点坐标是[0,0,0]
        无人机向左移动：Y轴减少数值
        无人机向右移动：Y轴增加数值
        无人机向上移动：Z轴减少数值
        无人机向下移动：Z轴增加数值
        无人机向后移动：X轴减少数值
        无人机向前移动：X轴增加数值'''
        import os

        # os.environ['CUDA_VISIBLE_DEVICES'] = '0'

        from swift.tuners import Swift

        ckpt_dir = '../llms/tuned_llms/qwen-1_8b-chat-int4/v0-20240327-221351/checkpoint-100'
        model_type = ModelType.qwen1half_1_8b_chat_int4
        template_type = get_default_template_type(model_type)
        print(template_type)
        self.model, tokenizer = get_model_tokenizer(model_type, model_dir="../llms/pretrained_llms/qwen/Qwen-1_8B-Chat-Int4", model_kwargs={'device_map': 'auto'})

        self.model = Swift.from_pretrained(self.model, ckpt_dir, inference_mode=True)
        self.template = get_template(template_type, tokenizer)

    def ask(self, question):
        response, history = inference(self.model, self.template, question, system=self.SystemPromots)
        return response


def extract_python_code(content):
    code_block_regex = re.compile(r"```(.*?)```", re.DOTALL)
    code_blocks = code_block_regex.findall(content)
    if code_blocks:
        full_code = "\n".join(code_blocks)

        if full_code.startswith("python"):
            full_code = full_code[7:]

        return full_code
    else:
        return None


class colors:  # You may need to change color settings
    RED = "\033[31m"
    ENDC = "\033[m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"

with open(args.prompt, "r") as f:
    prompt = f.read()

print("Welcome to the AirSim chatbot! I am ready to help you with your AirSim questions and commands.")

def main():
    if args.local_llm:
        print('use local llm')
        llm = LocalLlm()
    else:
        print('use cloud llm')
        llm = CloudLlm()

    llm.ask(prompt)
    
    while True:
        question = input(colors.YELLOW + "AirSim> " + colors.ENDC)

        if question == "!quit" or question == "!exit":
            break

        if question == "!clear":
            os.system("cls")
            continue

        response = llm.ask(question)

        # print(f"\n{response}\n")

        code = extract_python_code(response)

        if code is not None:
            print("response  =", response)
            print("Please wait while I run the code in AirSim...")
            exec(extract_python_code(response))
            print("Done!\n")

main()